import wheels.pi.aio as aio

from wheels.pi.aio.Proactor import Proactor, EventHandler
from wheels.pi.aio.AsyncOperation import AcceptOperation, ConnectOperation, RecvOperation, SendOperation



class AIOScheduler(object):

    class _EventHandler(EventHandler):

        def __init__(self, task_scheduler, tid):
            self._task_scheduler = task_scheduler
            self._tid = tid

        def handle_accept(self, client_socket, remote_address):
            self._task_scheduler.schedule_task(self._tid, val=(client_socket, remote_address))

        def handle_connect(self):
            self._task_scheduler.schedule_task(self._tid)

        def handle_recv(self, buf):
            self._task_scheduler.schedule_task(self._tid, val=buf)

        def handle_send(self, cnt):
            self._task_scheduler.schedule_task(self._tid, val=cnt)


    def __init__(self, task_scheduler):
        self._task_scheduler = task_scheduler
        self._proactor = Proactor()


    def schedule(self, tid, op):
        handler = AIOScheduler._EventHandler(self._task_scheduler, tid)
        aio.register(op, self._proactor, handler)


    def schedule_accept(self, tid, handle):
        op = AcceptOperation(handle)
        self.schedule(tid, op)


    def schedule_connect(self, tid, handle, remote_address):
        op = ConnectOperation(handle, remote_address)
        self.schedule(tid, op)


    def schedule_recv(self, tid, handle, cnt):
        op = RecvOperation(handle, cnt)
        self.schedule(tid, op)


    def schedule_send(self, tid, handle, buf):
        op = SendOperation(handle, buf)
        self.schedule(tid, op)


    def run(self):
        aio.flush()
